import simpy
import math
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from collections import namedtuple

from core.board import Board
from core.task import Task
from core.plane import Plane
from core.worker import Worker

from core.parameters import Parameters


class RLEnv:
    '''
        :Env rl交互环境
    '''

    def __init__(self, pa, worker_configs, plane_configs, dag_dataframe):
        self.env = None
        self.events = None
        self.board = None
        self.task_broker = None
        # self.scheduler =
        self.pa = pa
        self.plane_configs = plane_configs
        self.worker_configs = worker_configs
        # 设置plane的dag dataframe
        Plane.task_dataframe = dag_dataframe

        # self.task_broker.attach(self)
        self.board = Board()
        self.last_timestamp = -1
        # reset planes....
        for plane_config in self.plane_configs:
            # 静态保障
            # assert plane_config.submit_time >= self.env.now
            # 静态保障
            # yield self.env.timeout(0)
            plane = Plane(self.env, plane_config)
            # print('a plane arrived at time %f' % self.env.now)
            self.board.add_plane(plane)
        # reset workers...
        self.board.add_workers(self.worker_configs)
        self.end = ''
        self.error_count = 0
        Worker.time_horizon = pa.time_horizon

    @property
    def action_space(self) -> tuple:

        return self.board.num_group * self.pa.num_nw + 1
        # 0 ,1,2,...,self.pa.num_nw,...,self.board.num_group * self.pa.num_nw

    @property
    def observer_space(self) -> tuple:

        return (self.pa.time_horizon, self.board.num_group*4), \
               (self.pa.time_horizon, self.pa.num_nw*4), \
               (len(self.plane_configs), 4, 4)

    def reset(self):
        '''重置env'''
        self.env = simpy.Environment()
        self.events = [self.env.event()]
        self.last_timestamp = -1
        Task.event_init(self.events)
        # reset board
        self.board.reset(self.env)
        self.end = ''

        self.task_slot = TaskSlot(self.pa)
        # self.task_backlog = TaskBacklog(self.pa)
        i = 0
        task = self.board.next_task()

        Task.gantt = []  #
        # schedule_events = None
        Task.event_tick = -1

        self.error_count = 0
        self.error_move = False

        while task != None:
            if i < self.pa.num_nw:
                self.task_slot.slot[i] = task
            # elif i < self.pa.num_nw + self.pa.backlog_size:
            #     self.task_backlog.backlog[i - self.pa.num_nw] = task
            #     self.task_backlog.curr_size += 1

            i += 1
            if i == self.pa.num_nw:
                break
            task = self.board.next_task()

        return self.observe()

    def step(self, action: int) -> tuple:
        ''' 选着当前的action 并作出当前的env '''
        """
        while not env.done:
            
            action = RL(state)
            next_state,reward =env.step(action)
            # action 为0-M ,M为最大动作宽度，
            # 0和无效动作会执行仿真，其他动作只改变state,不执行仿真过程
            state = next_state
        """
        status = None

        done = False
        reward = 0
        info = None
        self.error_move = False
        # env.run(event) 设置当前的仿真事件停止
        task_act = action % self.pa.num_nw
        group_num = action // self.pa.num_nw

        if action == self.action_space:  # explicit void action
            status = 'MoveOn'
        elif self.task_slot.slot[task_act] is None:  # implicit void action
            status = 'MoveOn'
        else:
            allocated = self.allocate_task(group_num,task_act)
            if not allocated:  # implicit void action
                status = 'MoveOn'
            else:
                status = 'Allocate'

        if status == 'MoveOn':
            # simpy 运行到事件后自动停止时间戳的所有任务
            # todo:moveon的运行逻辑
            if self.env.now > self.pa.episode_max_length:  # run too long, force termination
                done = True

            # elif self.error_count >= 10:
            #
            #     done = True

            elif self.end == 'all_done':
                done = True
            elif self.end =='only_move':
                # move on to end
                if self.board.canMoveOn():
                    self.env.run(until=self.events[0])
                done = not(self.board.canMoveOn())
                self.end = 'all_done' if done else 'only_move'

            elif self.end == 'no_new_task':
                # 只需要完成当前的任务就可以了
                # task_backlog & task_slot
                done = False

            if not done:
                # 判断是否需要推进simpy env
                if self.board.canMoveOn():
                    self.env.run(until=self.events[0])
                else:
                    timestamp = self.env.now + 10
                    self.env.run(until=timestamp)
                    self.error_move = True
                    self.error_count += 1

                if self.board.has_task:
                    new_task = self.board.next_task()

                    if new_task is not None:  # a new job comes

                        to_board = True
                        for i in range(self.pa.num_nw):
                            if self.task_slot.slot[i] is None:  # put in new visible job slots
                                self.task_slot.slot[i] = new_task

                                to_board = False
                                break

                        if to_board:
                            self.board.insert_task(new_task)

                        # if to_backlog:
                            # if self.task_backlog.curr_size < self.pa.backlog_size:
                            #     self.task_backlog.backlog[self.task_backlog.curr_size] = new_task
                            #     self.task_backlog.curr_size += 1
                            #
                            # else:  # abort, backlog full
                            #     print("Backlog is full.")
                            #     self.board.insert_task(new_task)
                    else:
                        # 当前task 处于un_release状态，无法
                        self.end = 'has_un_released_task'

                else:
                    # if not self.board.canMoveOn() and\
                    #         all(s is None for s in self.task_slot.slot) and \
                    #         all(s is None for s in self.task_backlog.backlog):
                    #     self.end = 'all_done'
                    if all(s is None for s in self.task_slot.slot):# and \
                    #         self.task_backlog.curr_size == 0:
                        if self.board.canMoveOn():
                            self.end = 'only_move'
                            # 只需要运行当前worker完成任务
                        else:
                            self.end = 'all_done'
                            done = True
                    else:
                        self.end = 'no_new_task'

                                # exit(1)


                    # self.extra_info.new_job_comes()
            # if not done:

            reward = self.get_reward(done)
            self.last_timestamp = self.env.now

            # add new jobs
        elif status == 'Allocate':
            # self.task_record.record[self.job_slot.slot[action].id] = self.job_slot.slot[jb_idx]
            self.task_slot.slot[task_act] = None

            # dequeue backlog
            # if self.task_backlog.curr_size > 0:
            #     self.task_slot.slot[task_act] = self.task_backlog.backlog[0]  # if backlog empty, it will be 0
            #     self.task_backlog.backlog[: -1] = self.task_backlog.backlog[1:]  # move every job one unit size forward
            #     self.task_backlog.backlog[-1] = None
            #     self.task_backlog.curr_size -= 1

        ob = self.observe()

        info = None

        # if self.render:
        #     self.plot_state()

        return ob, reward, done, info

    def observe(self) -> tuple:
        '''

         根据self.task_slot,woker生成image

        '''

        time_horizon = self.pa.time_horizon
        worker_len = len(self.worker_configs)

        # worker_matrix = np.zeros((worker_len,time_horizon))
        workertype_dict = {'kp1':0,'kp2':1,'kp3':2,'kp4':3}

        worker_matrix = self.board.state

        slot_matrix = np.zeros((time_horizon, self.pa.num_nw * 4), dtype=np.float32)

        for i in range(self.pa.num_nw):
            task = self.task_slot.slot[i]
            if task is not None:
                slot_matrix[:task.duration,\
                  i*4+workertype_dict[task.worker_type]] = 1

        # backlog_matrix = np.zeros((self.pa.backlog_size,1))
        #
        # for i in range(self.task_backlog.curr_size):
        #     backlog_matrix[i,:] = 1

        # obs_mat = np.hstack((worker_matrix, slot_matrix))
        # # ,backlog_matrix))

        observes = namedtuple('observes','worker task graph')

        graph_matrix = self.board.task_state

        return observes(worker=worker_matrix, task=slot_matrix, graph=graph_matrix)

        # return obs_mat

    def get_reward(self,done):

        # if self.error_count >= 40:
        #     return -10 -3*(self.board.left_time)
        #
        # if self.last_timestamp == -1:
        #     return 0 if self.env.now == 0 else 1./self.env.now
        #
        # elif self.error_move:
        #     return -2

        if self.end == 'all_done' and done:
            return 45./self.env.now
        elif self.end == 'has_un_released_task' and done:
            return 45./self.env.now-self.board.left_time / 20.0 - self.left_task_time()/20.0
        elif self.end == 'no_new_task'and done:
            return 45./self.env.now-self.board.left_time / 20.0 - self.left_task_time()/20.0
        elif self.error_move:
            return -0.5
        return 0

    def left_task_time(self):
        duration = 0
        for task in self.task_slot.slot:
            if task is not None:
                duration += task.duration

        # for task in self.task_backlog.backlog:
        #     if task is not None:
        #         duration += task.duration
        return duration

    def allocate_task(self,group_num,task_act)-> bool:

        return self.board.allocate_task(group_num,self.task_slot.slot[task_act])

    def plot_state(self):
        plt.figure("screen", figsize=(20, 5))

        skip_row = 0

        for i in range(self.pa.num_res):

            plt.subplot(self.pa.num_res,
                        self.pa.num_machines + self.pa.num_nw + 1,
                        # 1 + self.pa.num_nw + 1,  # first +1 for current work, last +1 for backlog queue
                        i * (self.pa.num_nw + 1) + skip_row + 1)  # plot the backlog at the end, +1 to avoid 0

            # plt.imshow(self.machines.canvas[i, :, :], interpolation='nearest', vmax=1)

            # plt.show()

            for j in range(self.pa.num_nw):

                job_slot = np.zeros((self.pa.time_horizon, self.pa.max_job_size))
                if self.job_slot.slot[j] is not None:  # fill in a block of work
                    job_slot[: self.job_slot.slot[j].len, :self.job_slot.slot[j].res_vec[i]] = 1

                plt.subplot(self.pa.num_res,
                            1 + self.pa.num_nw + 1,  # first +1 for current work, last +1 for backlog queue
                            1 + i * (self.pa.num_nw + 1) + j + skip_row + 1)  # plot the backlog at the end, +1 to avoid 0

                plt.imshow(job_slot, interpolation='nearest', vmax=1)

                if j == self.pa.num_nw - 1:
                    skip_row += 1

        skip_row -= 1
        backlog_width = int(math.ceil(self.pa.backlog_size / float(self.pa.time_horizon)))
        backlog = np.zeros((self.pa.time_horizon, backlog_width))

        backlog[: self.job_backlog.curr_size / backlog_width, : backlog_width] = 1
        backlog[self.job_backlog.curr_size / backlog_width, : self.job_backlog.curr_size % backlog_width] = 1

        plt.subplot(self.pa.num_res,
                    1 + self.pa.num_nw + 1,  # first +1 for current work, last +1 for backlog queue
                    self.pa.num_nw + 1 + 1)

        plt.imshow(backlog, interpolation='nearest', vmax=1)

        plt.subplot(self.pa.num_res,
                    1 + self.pa.num_nw + 1,  # first +1 for current work, last +1 for backlog queue
                    self.pa.num_res * (self.pa.num_nw + 1) + skip_row + 1)  # plot the backlog at the end, +1 to avoid 0

        extra_info = np.ones((self.pa.time_horizon, 1)) * \
                     self.extra_info.time_since_last_new_job / \
                     float(self.extra_info.max_tracking_time_since_last_job)

        plt.imshow(extra_info, interpolation='nearest', vmax=1)

        plt.show()     # manual
        # plt.pause(0.01)  # automatic

    def plot_gantt(self):

        plt.figure(figsize=(24, 6))
        ax = plt.gca()
        [ax.spines[i].set＿visible(False) for i in ['top', 'right']]

        if self.end == 'all_done':
            gantt_list = Task.log_gantt()
            task_schedule = pd.DataFrame(data=gantt_list, columns=['task_id', 'task_type', 'worker_type', \
                                                                   'worker_id', 'start', 'end'])
            task_nums = len(task_schedule)
            alloc_workers_id = task_schedule.loc[:, 'worker_id']
            # workers = list(set(alloc_workers_id))
            # workers.sort()
            board_workers_id = self.board.get_workers_id
            workers_map = dict([(id, i) for i, id in enumerate(board_workers_id)])

            colors = ['#F0F8FF', '#F0F8FF', '#F0F8FF', '#FAEBD7', '#00FFFF', '#7FFFD4', '#F0FFFF', \
                      '#FFEBCD', '#0000FF', '#8A2BE2', '#A52A2A', '#DEB887', '#5F9EA0', 'm']

            color_map = dict()

            for i in range(ord('A'), ord('N') + 1):
                color_map[chr(i)] = colors[i - ord('A')]

            for idx in range(task_nums):
                task_type, worker_id, task_id = task_schedule.loc[idx, \
                                                                  ['task_type', 'worker_id', 'task_id']]
                # worker_id = task_schedule.loc[idx,'worker_id']
                y = workers_map[worker_id]
                left = task_schedule.loc[idx, 'start']
                width = task_schedule.loc[idx, 'end'] - left
                plt.barh(y, width=width, left=left, color=color_map[task_type], edgecolor='r')
                plt.text(left + 1 / 2, y, '%s' % task_id, \
                         color='r', size=7)

            plt.yticks(np.arange(len(board_workers_id)), np.arange(1, len(board_workers_id) + 1))
            # plt.show()
            plt.savefig(r'..\out\gantt.png')


class TaskSlot:
    '''task slot 存放task的表格'''

    def __init__(self, pa):
        self.slot = [None] * pa.num_nw


# class TaskBacklog:
#     '''存放task fifo的表，用于先放未进入task-slot的task-instances'''
#
#     def __init__(self, pa):
#         self.backlog = [None] * pa.backlog_size
#         self.curr_size = 0
#         self.max_size = pa.backlog_size
#
#     def isEmpty(self):
#         # if self.curr_size==0:
#         return self.curr_size == 0
#
#     def last(self):
#         if self.isEmpty():
#             return None
#         else:
#             return self.backlog[self.curr_size - 1]
#
#     def popfirst(self):
#         if self.isEmpty():
#             return None
#         else:
#             first = self.backlog[0]
#             for i in range(1, self.curr_size):
#                 self.backlog[i - 1] = self.backlog[i]
#             self.backlog[i] = None
#             self.curr_size -= 1
#
#     def isFull(self):
#         if self.max_size == self.curr_size:
#             return True
#         return False
#
#     def append(self, task):
#         if self.max_size == self.curr_size:
#             return False
#         else:
#             self.backlog[self.curr_size] = task
#             self.curr_size += 1


# ==========================================================================
# ------------------------------- Unit Tests -------------------------------
# ==========================================================================

# def test_backlog():
#     pa = Parameters()
#     pa.num_nw = 5
#     pa.simu_len = 50
#     pa.num_ex = 10
#     pa.new_job_rate = 1
#     # pa.compute_dependent_parameters()
#
#     env = RLEnv(pa, render=True, repre='image')
#     env.reset()
#
#     env.step(1)
#     env.step(1)
#     env.step(1)
#     env.step(1)
#     env.step(1)
#
#     env.step(5)
#     assert env.job_backlog.backlog[0] is not None
#     assert env.job_backlog.backlog[1] is None
#     print("New job is backlogged.")

# if __name__ == "__main__":
#
#     test_backlog()
